/**
 * @file ThreadPool.cpp
 * @brief 线程池实现
 * @author CloudLinker
 * @date 2023-06-14
 */

#include "CLink/Utils/ThreadPool.h"
#include <chrono>
#include <iostream>

namespace CLink
{
namespace Utils
{

void ThreadPool::Initialize(size_t threads)
{
    std::unique_lock<std::mutex> lock(mutex_);

    // 如果线程池已经初始化，则返回
    if (!workers_.empty())
    {
        return;
    }

    size_t hardware_threads = std::thread::hardware_concurrency();
    min_threads_ = std::min(min_threads_.load(), hardware_threads);
    max_threads_ = std::max(max_threads_.load(), hardware_threads);

    // 计算初始线程数
    size_t init_threads = threads > 0 ? threads : hardware_threads;
    init_threads = std::clamp(init_threads, min_threads_.load(), max_threads_.load());

    target_thread_count_ = init_threads;

    // 启动工作线程
    for (size_t i = 0; i < init_threads; ++i)
    {
        workers_.emplace_back(&ThreadPool::Worker, this);
    }

    // 启动监控线程
    std::thread monitor(&ThreadPool::MonitorThreads, this);
    monitor.detach();
}

ThreadPool::~ThreadPool()
{
    DrainAndStop();
}

// Worker线程主循环
void ThreadPool::Worker()
{
    while (!stop_)
    {
        PriorityTask task;

        {
            std::unique_lock<std::mutex> lock(mutex_);

            // 等待条件：停止或非暂停且任务可用
            condition_.wait(lock, [this]
                            { return stop_ || (!pause_ && !tasks_.empty()); });

            if (stop_)
                return;

            if (pause_)
                continue;

            task = std::move(tasks_.top());
            tasks_.pop();
            active_tasks_++;
        }

        // 执行任务
        task.function();

        // 通知可能的等待者
        if (drain_ && pending_tasks_ == 0)
        {
            drain_condition_.notify_one();
        }

        wait_condition_.notify_one();
    }
}

// 监控线程函数
void ThreadPool::MonitorThreads()
{
    while (!stop_)
    {
        std::this_thread::sleep_for(idle_thread_shrink_interval_);
        ShrinkIdleThreads();
        AdjustThreads();
    }
}

// 动态调整线程数
void ThreadPool::Resize(size_t threads)
{
    threads = std::clamp(threads, min_threads_.load(), max_threads_.load());

    {
        std::lock_guard<std::mutex> lock(mutex_);
        target_thread_count_ = threads;
        threads_adjusted_ = true;
    }

    condition_.notify_all();
}

// 增加空闲线程收缩
void ThreadPool::ShrinkIdleThreads()
{
    if (workers_.size() <= target_thread_count_)
        return;

    std::lock_guard<std::mutex> lock(mutex_);

    // 检查空闲线程
    size_t excess_threads = workers_.size() - target_thread_count_;
    if (excess_threads == 0)
        return;

    // 标记为停止并通知
    for (size_t i = 0; i < excess_threads; ++i)
    {
        condition_.notify_one();
    }

    // 等待并移除线程
    auto it = workers_.begin();
    while (it != workers_.end())
    {
        if (excess_threads == 0)
            break;

        if (it->joinable())
        {
            it->join();
            it = workers_.erase(it);
            excess_threads--;
        }
        else
        {
            ++it;
        }
    }
}

// 等待所有任务完成
void ThreadPool::WaitForAll()
{
    std::unique_lock<std::mutex> lock(mutex_);
    wait_condition_.wait(lock, [this]
                            { return pending_tasks_ == 0 && active_tasks_ == 0; });
}

// 带超时的等待
bool ThreadPool::WaitForAllFor(std::chrono::milliseconds timeout)
{
    std::unique_lock<std::mutex> lock(mutex_);
    return wait_condition_.wait_for(lock, timeout, [this]
                                    { return pending_tasks_ == 0 && active_tasks_ == 0; });
}

// 排空任务但不停止线程
void ThreadPool::Drain()
{
    drain_ = true;
    WaitForAll();
    drain_ = false;
}

// 带超时的排空
void ThreadPool::DrainFor(std::chrono::milliseconds timeout)
{
    drain_ = true;
    std::unique_lock<std::mutex> lock(mutex_);
    drain_condition_.wait_for(lock, timeout, [this]
                                { return pending_tasks_ == 0; });
    drain_ = false;
}

// 排空并停止
void ThreadPool::DrainAndStop()
{
    drain_ = true;
    stop_ = true;

    condition_.notify_all();

    for (std::thread &worker : workers_)
    {
        if (worker.joinable())
        {
            worker.join();
        }
    }

    workers_.clear();
}

// 暂停线程池
void ThreadPool::Pause()
{
    pause_ = true;
}

// 恢复线程池
void ThreadPool::Resume()
{
    pause_ = false;
    condition_.notify_all();
}

// 停止线程池
void ThreadPool::Stop(bool immediate)
{
    if (immediate)
    {
        stop_ = true;
        condition_.notify_all();
    }
    else
    {
        DrainAndStop();
    }
}

// 调整线程数
void ThreadPool::AdjustThreads()
{
    if (!threads_adjusted_)
        return;

    size_t target = target_thread_count_;
    size_t current = workers_.size();

    if (target > current)
    {
        std::lock_guard<std::mutex> lock(mutex_);
        for (size_t i = current; i < target; ++i)
        {
            workers_.emplace_back(&ThreadPool::Worker, this);
        }
    }
    else if (target < current)
    {
        ShrinkIdleThreads();
    }

    threads_adjusted_ = false;
}

// 获取活跃任务数
size_t ThreadPool::ActiveTasks() const
{
    return active_tasks_.load();
}

// 获取等待任务数
size_t ThreadPool::QueuedTasks() const
{
    std::lock_guard<std::mutex> lock(mutex_);
    return tasks_.size();
}

// 设置最小线程数
void ThreadPool::SetMinThreads(size_t min_threads)
{
    min_threads_.store(min_threads);
    min_threads_.store(std::min(min_threads_.load(), max_threads_.load()));
}

// 设置最大线程数
void ThreadPool::SetMaxThreads(size_t max_threads)
{
    max_threads_.store(max_threads);
    max_threads_.store(std::max(max_threads_.load(), min_threads_.load()));
}

// 获取当前线程数
size_t ThreadPool::TotalThreads() const
{
    std::lock_guard<std::mutex> lock(mutex_);
    return workers_.size();
}

// 其他getter/setter实现
size_t ThreadPool::GetMinThreads() const { return min_threads_; }
size_t ThreadPool::GetMaxThreads() const { return max_threads_; }
bool ThreadPool::IsPaused() const { return pause_; }
void ThreadPool::SetTaskCancellation(bool enable) { task_cancellation_ = enable; }
bool ThreadPool::IsTaskCancellationEnabled() const { return task_cancellation_; }
void ThreadPool::SetThreadBusyTimeout(std::chrono::milliseconds timeout) { thread_busy_timeout_ = timeout; }
std::chrono::milliseconds ThreadPool::GetThreadBusyTimeout() const { return thread_busy_timeout_; }
void ThreadPool::SetIdleThreadShrinkInterval(std::chrono::milliseconds interval) { idle_thread_shrink_interval_ = interval; }
std::chrono::milliseconds ThreadPool::GetIdleThreadShrinkInterval() const { return idle_thread_shrink_interval_; }

} // namespace Utils
} // namespace CLink